"""
双向认证(mTLS)模块 - GPU服务器与网关间安全通信

这个模块实现了双向TLS认证：
- 证书管理和验证
- 客户端和服务器双向认证
- 证书轮换支持
- 安全连接建立
- 身份验证和授权
"""

import asyncio
import logging
import ssl
import time
from dataclasses import dataclass, field
from datetime import datetime, timezone, timedelta
from enum import Enum
from typing import Any, Dict, List, Optional, Union, Tuple
import json
import hashlib
import base64
from pathlib import Path

import aiohttp
from cryptography import x509
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives.kdf import pbkdf2
from cryptography.x509.oid import NameOID, ExtendedKeyUsageOID
from cryptography.hazmat.primitives.serialization import pkcs12

logger = logging.getLogger(__name__)


class CertificateStatus(Enum):
    """证书状态"""
    VALID = "valid"
    EXPIRED = "expired"
    REVOKED = "revoked"
    INVALID = "invalid"
    PENDING = "pending"


class AuthRole(Enum):
    """认证角色"""
    CLIENT = "client"
    SERVER = "server"
    GATEWAY = "gateway"
    GPU_SERVER = "gpu_server"


@dataclass
class CertificateInfo:
    """证书信息"""
    subject: str
    issuer: str
    serial_number: str
    not_valid_before: datetime
    not_valid_after: datetime
    status: CertificateStatus
    fingerprint: str
    role: AuthRole
    metadata: Dict[str, Any] = field(default_factory=dict)


@dataclass
class MTLSConfig:
    """mTLS配置"""
    # 证书配置
    ca_cert_path: str = "certs/ca.crt"
    ca_key_path: str = "certs/ca.key"
    server_cert_path: str = "certs/server.crt"
    server_key_path: str = "certs/server.key"
    client_cert_path: str = "certs/client.crt"
    client_key_path: str = "certs/client.key"
    
    # 证书生成配置
    key_size: int = 2048
    cert_validity_days: int = 365
    country_name: str = "CN"
    state_name: str = "Beijing"
    city_name: str = "Beijing"
    organization_name: str = "GPU Auto Scaling System"
    
    # 安全配置
    cipher_suites: List[str] = field(default_factory=lambda: [
        "ECDHE-RSA-AES256-GCM-SHA384",
        "ECDHE-RSA-AES128-GCM-SHA256",
        "ECDHE-RSA-AES256-SHA384",
        "ECDHE-RSA-AES128-SHA256"
    ])
    
    # 验证配置
    verify_client_cert: bool = True
    verify_server_cert: bool = True
    require_client_auth: bool = True
    check_cert_revocation: bool = True
    
    # 连接配置
    connection_timeout: float = 10.0
    handshake_timeout: float = 5.0
    max_retries: int = 3


class CertificateManager:
    """证书管理器"""
    
    def __init__(self, config: Optional[MTLSConfig] = None):
        self.config = config or MTLSConfig()
        self.certificates: Dict[str, CertificateInfo] = {}
        self.revoked_certificates: set = set()
        
        # 确保证书目录存在
        self.cert_dir = Path("certs")
        self.cert_dir.mkdir(exist_ok=True)
        
        self.logger = logging.getLogger(__name__)
        self.logger.info("证书管理器已初始化")
    
    async def generate_ca_certificate(self) -> Tuple[str, str]:
        """生成CA证书"""
        try:
            # 生成私钥
            private_key = rsa.generate_private_key(
                public_exponent=65537,
                key_size=self.config.key_size
            )
            
            # 创建证书
            subject = issuer = x509.Name([
                x509.NameAttribute(NameOID.COUNTRY_NAME, self.config.country_name),
                x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, self.config.state_name),
                x509.NameAttribute(NameOID.LOCALITY_NAME, self.config.city_name),
                x509.NameAttribute(NameOID.ORGANIZATION_NAME, self.config.organization_name),
                x509.NameAttribute(NameOID.COMMON_NAME, "GPU Auto Scaling CA"),
            ])
            
            cert = x509.CertificateBuilder().subject_name(
                subject
            ).issuer_name(
                issuer
            ).public_key(
                private_key.public_key()
            ).serial_number(
                x509.random_serial_number()
            ).not_valid_before(
                datetime.utcnow()
            ).not_valid_after(
                datetime.utcnow() + timedelta(days=self.config.cert_validity_days)
            ).add_extension(
                x509.BasicConstraints(ca=True, path_length=None),
                critical=True,
            ).add_extension(
                x509.KeyUsage(
                    key_cert_sign=True,
                    crl_sign=True,
                    digital_signature=True,
                    key_encipherment=False,
                    content_commitment=False,
                    data_encipherment=False,
                    key_agreement=False,
                    encipher_only=False,
                    decipher_only=False
                ),
                critical=True,
            ).sign(private_key, hashes.SHA256())
            
            # 保存证书和私钥
            ca_cert_path = self.cert_dir / "ca.crt"
            ca_key_path = self.cert_dir / "ca.key"
            
            with open(ca_cert_path, "wb") as f:
                f.write(cert.public_bytes(serialization.Encoding.PEM))
            
            with open(ca_key_path, "wb") as f:
                f.write(private_key.private_bytes(
                    encoding=serialization.Encoding.PEM,
                    format=serialization.PrivateFormat.PKCS8,
                    encryption_algorithm=serialization.NoEncryption()
                ))
            
            self.logger.info("CA证书生成成功")
            return str(ca_cert_path), str(ca_key_path)
            
        except Exception as e:
            self.logger.error(f"CA证书生成失败: {e}")
            raise
    
    async def generate_server_certificate(self, 
                                        common_name: str,
                                        san_dns: Optional[List[str]] = None,
                                        san_ip: Optional[List[str]] = None) -> Tuple[str, str]:
        """生成服务器证书"""
        try:
            # 生成私钥
            private_key = rsa.generate_private_key(
                public_exponent=65537,
                key_size=self.config.key_size
            )
            
            # 创建证书
            subject = x509.Name([
                x509.NameAttribute(NameOID.COUNTRY_NAME, self.config.country_name),
                x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, self.config.state_name),
                x509.NameAttribute(NameOID.LOCALITY_NAME, self.config.city_name),
                x509.NameAttribute(NameOID.ORGANIZATION_NAME, self.config.organization_name),
                x509.NameAttribute(NameOID.COMMON_NAME, common_name),
            ])
            
            # 加载CA证书和私钥
            ca_cert, ca_key = await self._load_ca_certificate()
            
            cert_builder = x509.CertificateBuilder().subject_name(
                subject
            ).issuer_name(
                ca_cert.subject
            ).public_key(
                private_key.public_key()
            ).serial_number(
                x509.random_serial_number()
            ).not_valid_before(
                datetime.utcnow()
            ).not_valid_after(
                datetime.utcnow() + timedelta(days=self.config.cert_validity_days)
            )
            
            # 添加SAN扩展
            if san_dns or san_ip:
                san_list = []
                if san_dns:
                    san_list.extend([x509.DNSName(name) for name in san_dns])
                if san_ip:
                    san_list.extend([x509.IPAddress(ip) for ip in san_ip])
                
                cert_builder = cert_builder.add_extension(
                    x509.SubjectAlternativeName(san_list),
                    critical=False,
                )
            
            # 添加扩展密钥用法
            cert_builder = cert_builder.add_extension(
                x509.ExtendedKeyUsage([
                    ExtendedKeyUsageOID.SERVER_AUTH,
                    ExtendedKeyUsageOID.CLIENT_AUTH
                ]),
                critical=True,
            )
            
            cert = cert_builder.sign(ca_key, hashes.SHA256())
            
            # 保存证书和私钥
            server_cert_path = self.cert_dir / f"{common_name}_server.crt"
            server_key_path = self.cert_dir / f"{common_name}_server.key"
            
            with open(server_cert_path, "wb") as f:
                f.write(cert.public_bytes(serialization.Encoding.PEM))
            
            with open(server_key_path, "wb") as f:
                f.write(private_key.private_bytes(
                    encoding=serialization.Encoding.PEM,
                    format=serialization.PrivateFormat.PKCS8,
                    encryption_algorithm=serialization.NoEncryption()
                ))
            
            self.logger.info(f"服务器证书生成成功: {common_name}")
            return str(server_cert_path), str(server_key_path)
            
        except Exception as e:
            self.logger.error(f"服务器证书生成失败: {e}")
            raise
    
    async def generate_client_certificate(self, 
                                        common_name: str,
                                        role: AuthRole = AuthRole.CLIENT) -> Tuple[str, str]:
        """生成客户端证书"""
        try:
            # 生成私钥
            private_key = rsa.generate_private_key(
                public_exponent=65537,
                key_size=self.config.key_size
            )
            
            # 创建证书
            subject = x509.Name([
                x509.NameAttribute(NameOID.COUNTRY_NAME, self.config.country_name),
                x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, self.config.state_name),
                x509.NameAttribute(NameOID.LOCALITY_NAME, self.config.city_name),
                x509.NameAttribute(NameOID.ORGANIZATION_NAME, self.config.organization_name),
                x509.NameAttribute(NameOID.COMMON_NAME, common_name),
            ])
            
            # 加载CA证书和私钥
            ca_cert, ca_key = await self._load_ca_certificate()
            
            cert = x509.CertificateBuilder().subject_name(
                subject
            ).issuer_name(
                ca_cert.subject
            ).public_key(
                private_key.public_key()
            ).serial_number(
                x509.random_serial_number()
            ).not_valid_before(
                datetime.utcnow()
            ).not_valid_after(
                datetime.utcnow() + timedelta(days=self.config.cert_validity_days)
            ).add_extension(
                x509.ExtendedKeyUsage([
                    ExtendedKeyUsageOID.CLIENT_AUTH,
                    ExtendedKeyUsageOID.SERVER_AUTH
                ]),
                critical=True,
            ).sign(ca_key, hashes.SHA256())
            
            # 保存证书和私钥
            client_cert_path = self.cert_dir / f"{common_name}_client.crt"
            client_key_path = self.cert_dir / f"{common_name}_client.key"
            
            with open(client_cert_path, "wb") as f:
                f.write(cert.public_bytes(serialization.Encoding.PEM))
            
            with open(client_key_path, "wb") as f:
                f.write(private_key.private_bytes(
                    encoding=serialization.Encoding.PEM,
                    format=serialization.PrivateFormat.PKCS8,
                    encryption_algorithm=serialization.NoEncryption()
                ))
            
            # 记录证书信息
            cert_info = await self._parse_certificate_info(cert, role)
            self.certificates[common_name] = cert_info
            
            self.logger.info(f"客户端证书生成成功: {common_name}")
            return str(client_cert_path), str(client_key_path)
            
        except Exception as e:
            self.logger.error(f"客户端证书生成失败: {e}")
            raise
    
    async def _load_ca_certificate(self) -> Tuple[x509.Certificate, Any]:
        """加载CA证书和私钥"""
        ca_cert_path = Path(self.config.ca_cert_path)
        ca_key_path = Path(self.config.ca_key_path)
        
        if not ca_cert_path.exists() or not ca_key_path.exists():
            await self.generate_ca_certificate()
        
        # 加载证书
        with open(ca_cert_path, "rb") as f:
            ca_cert = x509.load_pem_x509_certificate(f.read())
        
        # 加载私钥
        with open(ca_key_path, "rb") as f:
            ca_key = serialization.load_pem_private_key(f.read(), password=None)
        
        return ca_cert, ca_key
    
    async def _parse_certificate_info(self, 
                                    cert: x509.Certificate,
                                    role: AuthRole) -> CertificateInfo:
        """解析证书信息"""
        # 计算指纹
        cert_bytes = cert.public_bytes(serialization.Encoding.DER)
        fingerprint = hashlib.sha256(cert_bytes).hexdigest()
        
        return CertificateInfo(
            subject=str(cert.subject),
            issuer=str(cert.issuer),
            serial_number=str(cert.serial_number),
            not_valid_before=cert.not_valid_before,
            not_valid_after=cert.not_valid_after,
            status=CertificateStatus.VALID,
            fingerprint=fingerprint,
            role=role
        )
    
    async def verify_certificate(self, cert_path: str) -> CertificateInfo:
        """验证证书"""
        try:
            with open(cert_path, "rb") as f:
                cert = x509.load_pem_x509_certificate(f.read())
            
            # 检查证书是否过期
            now = datetime.utcnow()
            if now < cert.not_valid_before or now > cert.not_valid_after:
                status = CertificateStatus.EXPIRED
            else:
                status = CertificateStatus.VALID
            
            # 检查证书是否被撤销
            serial_number = str(cert.serial_number)
            if serial_number in self.revoked_certificates:
                status = CertificateStatus.REVOKED
            
            cert_info = await self._parse_certificate_info(cert, AuthRole.CLIENT)
            cert_info.status = status
            
            return cert_info
            
        except Exception as e:
            self.logger.error(f"证书验证失败: {e}")
            raise
    
    async def revoke_certificate(self, serial_number: str) -> bool:
        """撤销证书"""
        try:
            self.revoked_certificates.add(serial_number)
            self.logger.info(f"证书已撤销: {serial_number}")
            return True
        except Exception as e:
            self.logger.error(f"证书撤销失败: {e}")
            return False


class MTLSClient:
    """mTLS客户端"""
    
    def __init__(self, 
                 cert_manager: CertificateManager,
                 config: Optional[MTLSConfig] = None):
        self.cert_manager = cert_manager
        self.config = config or MTLSConfig()
        self.session: Optional[aiohttp.ClientSession] = None
        
        self.logger = logging.getLogger(__name__)
    
    async def create_ssl_context(self, 
                                cert_path: str,
                                key_path: str,
                                ca_cert_path: str) -> ssl.SSLContext:
        """创建SSL上下文"""
        try:
            context = ssl.create_default_context(ssl.Purpose.SERVER_AUTH)
            context.load_cert_chain(cert_path, key_path)
            context.load_verify_locations(ca_cert_path)
            context.check_hostname = False  # 在mTLS中通常不检查主机名
            context.verify_mode = ssl.CERT_REQUIRED
            
            # 设置密码套件
            if hasattr(context, 'set_ciphers'):
                context.set_ciphers(':'.join(self.config.cipher_suites))
            
            return context
            
        except Exception as e:
            self.logger.error(f"SSL上下文创建失败: {e}")
            raise
    
    async def create_session(self, 
                           cert_path: str,
                           key_path: str,
                           ca_cert_path: str) -> aiohttp.ClientSession:
        """创建带mTLS的会话"""
        try:
            ssl_context = await self.create_ssl_context(cert_path, key_path, ca_cert_path)
            
            connector = aiohttp.TCPConnector(ssl=ssl_context)
            timeout = aiohttp.ClientTimeout(total=self.config.connection_timeout)
            
            session = aiohttp.ClientSession(
                connector=connector,
                timeout=timeout
            )
            
            self.session = session
            self.logger.info("mTLS客户端会话创建成功")
            return session
            
        except Exception as e:
            self.logger.error(f"mTLS客户端会话创建失败: {e}")
            raise
    
    async def request(self, 
                     method: str,
                     url: str,
                     **kwargs) -> aiohttp.ClientResponse:
        """发送mTLS请求"""
        if not self.session:
            raise RuntimeError("客户端会话未创建")
        
        try:
            response = await self.session.request(method, url, **kwargs)
            return response
        except Exception as e:
            self.logger.error(f"mTLS请求失败: {e}")
            raise
    
    async def close(self) -> None:
        """关闭客户端会话"""
        if self.session:
            await self.session.close()
            self.session = None


class MTLSServer:
    """mTLS服务器"""
    
    def __init__(self, 
                 cert_manager: CertificateManager,
                 config: Optional[MTLSConfig] = None):
        self.cert_manager = cert_manager
        self.config = config or MTLSConfig()
        
        self.logger = logging.getLogger(__name__)
    
    async def create_ssl_context(self, 
                                cert_path: str,
                                key_path: str,
                                ca_cert_path: str) -> ssl.SSLContext:
        """创建服务器SSL上下文"""
        try:
            context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
            context.load_cert_chain(cert_path, key_path)
            context.load_verify_locations(ca_cert_path)
            
            if self.config.require_client_auth:
                context.verify_mode = ssl.CERT_REQUIRED
            else:
                context.verify_mode = ssl.CERT_OPTIONAL
            
            # 设置密码套件
            if hasattr(context, 'set_ciphers'):
                context.set_ciphers(':'.join(self.config.cipher_suites))
            
            return context
            
        except Exception as e:
            self.logger.error(f"服务器SSL上下文创建失败: {e}")
            raise


class MTLSManager:
    """mTLS管理器"""
    
    def __init__(self, config: Optional[MTLSConfig] = None):
        self.config = config or MTLSConfig()
        self.cert_manager = CertificateManager(self.config)
        self.clients: Dict[str, MTLSClient] = {}
        self.servers: Dict[str, MTLSServer] = {}
        
        self.logger = logging.getLogger(__name__)
        self.logger.info("mTLS管理器已初始化")
    
    async def initialize_certificates(self) -> None:
        """初始化证书"""
        try:
            # 生成CA证书（如果不存在）
            ca_cert_path = Path(self.config.ca_cert_path)
            if not ca_cert_path.exists():
                await self.cert_manager.generate_ca_certificate()
            
            # 生成服务器证书（如果不存在）
            server_cert_path = Path(self.config.server_cert_path)
            if not server_cert_path.exists():
                await self.cert_manager.generate_server_certificate(
                    "gpu-auto-scaling-server",
                    san_dns=["localhost", "gpu-server"],
                    san_ip=["127.0.0.1", "0.0.0.0"]
                )
            
            # 生成客户端证书（如果不存在）
            client_cert_path = Path(self.config.client_cert_path)
            if not client_cert_path.exists():
                await self.cert_manager.generate_client_certificate(
                    "gpu-auto-scaling-client",
                    AuthRole.CLIENT
                )
            
            self.logger.info("证书初始化完成")
            
        except Exception as e:
            self.logger.error(f"证书初始化失败: {e}")
            raise
    
    async def create_client(self, 
                           name: str,
                           cert_path: str,
                           key_path: str,
                           ca_cert_path: str) -> MTLSClient:
        """创建mTLS客户端"""
        try:
            client = MTLSClient(self.cert_manager, self.config)
            await client.create_session(cert_path, key_path, ca_cert_path)
            
            self.clients[name] = client
            self.logger.info(f"mTLS客户端创建成功: {name}")
            return client
            
        except Exception as e:
            self.logger.error(f"mTLS客户端创建失败: {name}, 错误: {e}")
            raise
    
    async def create_server(self, 
                           name: str,
                           cert_path: str,
                           key_path: str,
                           ca_cert_path: str) -> MTLSServer:
        """创建mTLS服务器"""
        try:
            server = MTLSServer(self.cert_manager, self.config)
            ssl_context = await server.create_ssl_context(cert_path, key_path, ca_cert_path)
            
            self.servers[name] = server
            self.logger.info(f"mTLS服务器创建成功: {name}")
            return server
            
        except Exception as e:
            self.logger.error(f"mTLS服务器创建失败: {name}, 错误: {e}")
            raise
    
    async def get_certificate_status(self) -> Dict[str, Any]:
        """获取证书状态"""
        status = {
            "certificates": {},
            "revoked_certificates": list(self.cert_manager.revoked_certificates),
            "clients": list(self.clients.keys()),
            "servers": list(self.servers.keys())
        }
        
        for name, cert_info in self.cert_manager.certificates.items():
            status["certificates"][name] = {
                "subject": cert_info.subject,
                "issuer": cert_info.issuer,
                "serial_number": cert_info.serial_number,
                "not_valid_before": cert_info.not_valid_before.isoformat(),
                "not_valid_after": cert_info.not_valid_after.isoformat(),
                "status": cert_info.status.value,
                "fingerprint": cert_info.fingerprint,
                "role": cert_info.role.value
            }
        
        return status
    
    async def close_all_clients(self) -> None:
        """关闭所有客户端"""
        for client in self.clients.values():
            await client.close()
        self.clients.clear()
    
    async def revoke_certificate(self, serial_number: str) -> bool:
        """撤销证书"""
        return await self.cert_manager.revoke_certificate(serial_number)


# 全局mTLS管理器
_global_mtls_manager: Optional[MTLSManager] = None


def get_mtls_manager() -> MTLSManager:
    """获取全局mTLS管理器"""
    global _global_mtls_manager
    if _global_mtls_manager is None:
        _global_mtls_manager = MTLSManager()
    return _global_mtls_manager
